#!/usr/bin/env python
# coding: utf-8

# In[1]:


#加载必要的库
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torchvision import datasets,transforms

#定义超参数
BATCH_SIZE = 16  # 每批处理的数据
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu") # 是否用GPU还是CPU训练

import torch

# 检查是否有CUDA支持
if torch.cuda.is_available():
    # 加载模型时将存储映射到CPU
    saved_model_weights = torch.load('trained_model.pth', map_location=torch.device('cpu'))
else:
    # 正常加载模型
    saved_model_weights = torch.load('trained_model.pth')



EPOCHS = 15# 训练数据集的轮次
#  构建pipeline,对图像做处理
pipeline = transforms.Compose([
    transforms.ToTensor(),# 将图片转换成tensor
    transforms.Normalize((0.1307,),(0.3081,)) # 正则化降低模型复杂度
])
#下载、加载数据
from torch.utils.data import DataLoader

# 下载数据集
pipeline = transforms.Compose([
    transforms.Grayscale(num_output_channels=1),
    transforms.Resize((28, 28)),  # 将图像大小调整为 28x28
    transforms.ToTensor(),# 将图片转换成tensor
    transforms.Normalize((0.1307,),(0.3081,)) # 正则化降低模型复杂度
])
# 加载完整的训练集
train_dataset = datasets.ImageFolder('mnist+', transform=pipeline)
# 定义训练集和测试集的比例
train_ratio = 0.8  # 训练集占总体的80%
test_ratio = 0.2   # 测试集占总体的20%
# 计算划分的大小
train_size = int(train_ratio * len(train_dataset))
test_size = len(train_dataset) - train_size
# 使用random_split函数进行划分
train_dataset, test_dataset = torch.utils.data.random_split(train_dataset, [train_size, test_size])


# In[3]:


# 创建训练集和测试集的数据加载器
batch_size = 16
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=batch_size, shuffle=True)
#  构建网络模型
class Digit(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(1, 16, 5)  # 1:灰度图片的通道， 16：输出通道， 5：kernel 5x5
        self.conv2 = nn.Conv2d(16, 20, 3)  # 16:输入通道， 20：输出通道， 3：kernel 3x3
        self.fc1 = nn.Linear(2000, 500)  # 20*10*10:输入通道， 500：输出通道
        self.fc2 = nn.Linear(500, 16)  # 500：输入通道， 10：输出通道

    # 前向传播
    def forward(self, x):
        input_size = x.size(0)  # batch_size
        x = self.conv1(x)  # 输入：batch*1*28*28, 输出：batch*10*24*24 (28 - 5 + 1 = 24)
        x = F.relu(x)  # 激活函数，保持shape不变， 输出batch*10*24*24
        x = F.max_pool2d(x, 2, 2)  # 输入：batch*10*24*24 输出：batch*10*12*12

        x = self.conv2(x)  # 输入：batch*10*12*12, 输出：batch*20*10*10 (12 - 3 + 1 = 10)
        x = F.relu(x)

        x = x.view(input_size, -1)  # 拉平， -1 自动计算维度，20*10*10 = 2000

        x = self.fc1(x)  # 输入：batch*2000，输出：batch*500
        x = F.relu(x)

        x = self.fc2(x)  # 输入：batch*500，输出：batch*10

        output = F.log_softmax(x, dim=1)  # 计算分类后，每个数字的概率值

        return output


#  定义训练方法
def train_model(model, device, train_loader, optimizer, epoch):
    # 模型训练
    model.train()
    for batch_index, (data, target) in enumerate(train_loader):
        # 部署到DEVICE上去
        data, target = data.to(device), target.to(device)
        # 梯度初始化为0
        optimizer.zero_grad()
        # 训练后的结果
        output = model(data)
        # 计算损失
        loss = F.cross_entropy(output, target)  # 交叉熵用于分类比较多的情况

        # 反向传播
        loss.backward()
        # 参数优化
        optimizer.step()
        if batch_index % 4800 == 0:
            print("Train Epoch : {} \t Loss : {:.6f}".format(epoch, loss.item()))
# 定义测试方法
def test_model(model, device, test_loader):
    # 模型验证
    model.eval()
    # 正确率
    correct = 0.0
    # 测试损失
    test_loss = 0.0
    with torch.no_grad(): # 不会计算梯度，也不会反向传播
        for data, target in test_loader:
            # 部署到device上
            data, target = data.to(device), target.to(device)
            # 测试数据
            output = model(data)
            # 计算测试损失
            test_loss += F.cross_entropy(output, target).item()
            # 找到概率值最大的下标
            pred = output.max(1, keepdim=True)[1] # 值，索引
            # pred = torch.max(output, dim=1)
            # pred = output.argmax(dim=1)
            #累计正确的值
            correct += pred.eq(target.view_as(pred)).sum().item()
        test_loss /= len(test_loader.dataset)
        print("Test — Average loss : {:.4f}, Accuracy : {:.3f}\n".format(
                test_loss, 100.0 * correct / len(test_loader.dataset)))
model = Digit().to(DEVICE)
optimizer = optim.Adam(model.parameters())
for epoch in range(1, EPOCHS):
    train_model(model, DEVICE, train_loader, optimizer, epoch)
    # 保存训练后的模型权重
test_model(model, DEVICE, test_loader)
torch.save(model.state_dict(), 'trained_model.pth')


# In[ ]:




